package state;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.util.ArrayList;
import java.util.List;

public class OperatorStateDemo implements SinkFunction<Tuple2<Long, Long>>, CheckpointedFunction {

    ListState<Tuple2<Long, Long>> operatorState;
    int threshold;

    private List<Tuple2<Long, Long>> buffer;

    OperatorStateDemo(int threshold) {
        this.threshold = threshold;
        this.buffer = new ArrayList<>();
    }

    @Override
    public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception {
        System.out.println("....snapshotState");
        this.operatorState.clear();
        for(Tuple2<Long, Long> element: buffer) {
            operatorState.add(element);
        }
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        System.out.println("...initializeState");
        ListStateDescriptor<Tuple2<Long, Long>> descriptor = new ListStateDescriptor<>("operatorDemo", TypeInformation.of(new TypeHint<Tuple2<Long, Long>>() {
        }));
        operatorState = context.getOperatorStateStore().getListState(descriptor);
        if (context.isRestored()) {
            for (Tuple2<Long, Long> element: operatorState.get()) {
                buffer.add(element);
            }
            System.out.println("...context.isRestored(): true" + buffer);
        }
    }

    @Override
    public void invoke(Tuple2<Long, Long> value, Context context) throws Exception {
        System.out.println("---------incoke------------");
        buffer.add(value);
        if (buffer.size() == threshold) {
            for (Tuple2<Long, Long> element: buffer) {
                System.out.println("...out:" + element);
            }
            buffer.clear();
        }

    }
}
